import cv2 as cv
import matplotlib.pyplot as plt
img1 = cv.imread('assets/book.jpg', cv.IMREAD_GRAYSCALE)
img2 = cv.imread('assets/book3.jpg', cv.IMREAD_GRAYSCALE)

orb = cv.ORB_create()

kp1, des1 = orb.detectAndCompute(img1, None)
kp2, des2 = orb.detectAndCompute(img2, None)

# FLANN parameters
FLANN_INDEX_LSH = 6
index_params= dict(algorithm = FLANN_INDEX_LSH,
                   table_number = 12, # 12
                   key_size = 20,     # 20
                   multi_probe_level = 2) #2
search_params = dict(checks=30)
flann = cv.FlannBasedMatcher(index_params, search_params)
matches = flann.knnMatch(des1,des2,k=2)
print(len(matches))

matchesMask = [[0,0] for i in range(len(matches))]
# ratio test as per Lowe's paper
count=0
for i, p in enumerate(matches):
    if len(p) != 2:
        continue
    m, n = p
    if m.distance < 0.70*n.distance:
        matchesMask[i]=[1,0]
        count+=1
print(count)

draw_params = dict(matchColor = (0,255,0),
                   singlePointColor = (255,0,0),
                   matchesMask = matchesMask,
                   flags = cv.DrawMatchesFlags_DEFAULT)

img3 = cv.drawMatchesKnn(img1, kp1, img2, kp2, matches, None, **draw_params)
plt.imshow(img3)
plt.show()